#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient, ASCENDING, DESCENDING
import re
import datetime
import sys

default_encoding = 'utf-8'
if sys.getdefaultencoding() != default_encoding:
    reload(sys)
    sys.setdefaultencoding(default_encoding)


# 连接MongoDB
def InitMongoDB():
    MONGODB_HOST = 'localhost'
    MONGODB_PORT = 27017
    MONGODB_DB_NAME = 'Dictionary'
    MONGODB_COLLECTION_NAME = 'GeneOntology_20170504'
    connection = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = connection[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return connection, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def ParseOBO(filepath, lineSymbol, MongoDB_Collection):
    infile = open(filepath, "r")
    current_title = ""
    current_word = ""
    for line in infile:
        line = line.strip()
        if line.startswith("[") and line.endswith("]"):
            if current_title == "Term":
                insert_record = ParseGOTerm(current_word[:-len(lineSymbol)], lineSymbol)
                if len(insert_record) != 0:
                    MongoDB_Collection.insert_one(insert_record)
            current_title = line[1:-1]
            current_word = ""
        elif line != "":
            current_word += line + lineSymbol
    infile.close()
    insert_record = ParseGOTerm(current_word[:-len(lineSymbol)], lineSymbol)
    if len(insert_record) != 0:
        MongoDB_Collection.insert_one(insert_record)


def ParseGOTerm(term_word, lineSymbol):
    insert_record = {}
    relation_title = set(
        ["is_a", "intersection_of", "union_of", "disjoint_from", "relationship", "replaced_by", "consider"])
    if term_word != "":
        splitline = term_word.split(lineSymbol)

        for line in splitline:
            line = line.strip()
            if line.startswith("id:"):
                id = line.split(":")[-1].strip()
                insert_record["id"] = id
                insert_record["dms_id"] = id
                insert_record["dms_ids"] = [{"db": "GO", "id": id}]
                insert_record["dms_ids_extend"] = [{"db": "GO", "id": id}]
            elif line.startswith("name:"):
                name = line.split(":")[-1].strip()
                insert_record["name"] = name
                insert_record["dms_name"] = name
                insert_record["dms_synonym"] = [name]
                insert_record["dms_synonym_extend"] = [name]
            elif line.startswith("namespace:"):
                namespace = line.split(":")[-1].strip()
                insert_record["namespace"] = namespace
            elif line.startswith("alt_id:"):
                altid = line.split(":")[-1].strip()
                if insert_record.has_key("alt_id"):
                    insert_record["alt_id"].append(altid)
                else:
                    insert_record["alt_id"] = [altid]
                if insert_record.has_key("dms_ids"):
                    insert_record["dms_ids"].append({"db": "GO", "id": altid})
                else:
                    insert_record["dms_ids"] = [{"db": "GO", "id": altid}]
                if insert_record.has_key("dms_ids_extend"):
                    insert_record["dms_ids_extend"].append({"db": "GO", "id": altid})
                else:
                    insert_record["dms_ids_extend"] = [{"db": "GO", "id": altid}]
            elif line.startswith("def:"):
                def_def, def_other = getDefMatch(line)
                if def_def != None and def_other != None:
                    insert_record["def"] = def_def + "%split%" + def_other
                else:
                    print "[ERROR DEF] " + line
            elif line.startswith("comment:"):
                comment = line.split(":")[-1].strip()
                insert_record["comment"] = comment
            elif line.startswith("subset:"):
                subset = line.split(":")[-1].strip()
                if insert_record.has_key("subset"):
                    insert_record["subset"].append(subset)
                else:
                    insert_record["subset"] = [subset]
            elif line.startswith("synonym:"):
                syn_syn, syn_other = getSynonymMatch(line)
                if syn_syn != None and syn_other != None:
                    if insert_record.has_key("synonym"):
                        insert_record["synonym"].append(syn_syn + "%split%" + syn_other)
                    else:
                        insert_record["synonym"] = [syn_syn + "%split%" + syn_other]
                    if insert_record.has_key("dms_synonym"):
                        insert_record["dms_synonym"].append(syn_syn)
                    else:
                        insert_record["dms_synonym"] = [syn_syn]
                    if insert_record.has_key("dms_synonym_extend"):
                        insert_record["dms_synonym_extend"].append(syn_syn)
                    else:
                        insert_record["dms_synonym_extend"] = [syn_syn]
                else:
                    print "[ERROR SYNONYM] " + line
            elif line.startswith("xref:"):
                xrefs = line.split(":")
                xref_db = xrefs[1].strip()
                xref_id = xrefs[2].strip()
                if insert_record.has_key("xref"):
                    insert_record["xref"].append({"db": xref_db, "id": xref_id})
                else:
                    insert_record["xref"] = [{"db": xref_db, "id": xref_id}]
                if insert_record.has_key("dms_ids"):
                    insert_record["dms_ids"].append({"db": xref_db, "id": xref_id})
                else:
                    insert_record["dms_ids"] = [{"db": xref_db, "id": xref_id}]
                if insert_record.has_key("dms_ids_extend"):
                    insert_record["dms_ids_extend"].append({"db": xref_db, "id": xref_id})
                else:
                    insert_record["dms_ids_extend"] = [{"db": xref_db, "id": xref_id}]
            elif line.startswith("is_obsolete"):
                insert_record["is_obsolete"] = "true"
            elif line.startswith("created_by:"):
                creator = line.split(":")[1].strip()
                insert_record["created_by"] = creator
            elif line.startswith("creation_date:"):
                creatime = datetime.datetime.strptime(line[15:].strip(), "%Y-%m-%dT%H:%M:%SZ")
                insert_record["creation_date:"] = creatime
            elif line.startswith("property_value:"):
                propertyvalue = line[16:]
                if insert_record.has_key("property_value"):
                    insert_record["property_value"].append(propertyvalue)
                else:
                    insert_record["property_value"] = [propertyvalue]
            else:
                linetitle = line.split(":")[0].strip()
                if linetitle in relation_title:
                    relation_attr, relation_db, relation_id = getRelationMatch(line, linetitle)
                    if relation_db != None and relation_id != None:
                        intersection_insert = {"db": relation_db, "id": relation_id}
                        if relation_attr != None:
                            intersection_insert["attr"] = relation_attr
                        if insert_record.has_key(linetitle):
                            insert_record[linetitle].append(intersection_insert)
                        else:
                            insert_record[linetitle] = [intersection_insert]
                    else:
                        print "[ERROR " + linetitle.upper() + "] " + line
                else:
                    print "[ERROR UNKNOWN] " + line

    return insert_record


def getDefMatch(line):
    # regex = '^def: "(.*)" (\[.*\]) *$'
    regex = '^def: "(.*)" (\[.*\].*)$'
    match = re.match(regex, line)
    if match:
        rs1 = match.group(1).strip()
        rs2 = match.group(2).strip()
    else:
        rs1 = None
        rs2 = None
    return rs1, rs2


def getSynonymMatch(line):
    regex = '^synonym: "(.*)" ((EXACT|BROAD|NARROW|RELATED).*\[.*\].*)$'
    match = re.match(regex, line)
    if match:
        rs1 = match.group(1).strip()
        rs2 = match.group(2).strip()
    else:
        rs1 = None
        rs2 = None
    return rs1, rs2


def getRelationMatch(line, symbol):
    regex = '^' + symbol + ': (.* )?(.+):(.+) ! .*$'
    match = re.match(regex, line)
    if match:
        rs1 = match.group(1)
        if rs1 != None:
            rs1 = rs1.strip()
        rs2 = match.group(2).strip()
        rs3 = match.group(3).strip()
    else:
        regex2 = '^' + symbol + ': (.* )?(.+):(.+)$'
        match2 = re.match(regex2, line)
        if match2:
            rs1 = match2.group(1)
            if rs1 != None:
                rs1 = rs1.strip()
            rs2 = match2.group(2).strip()
            rs3 = match2.group(3).strip()
        else:
            rs1 = None
            rs2 = None
            rs3 = None
    return rs1, rs2, rs3


if __name__ == "__main__":
    MongoDBConnection, MongoDBCollection = InitMongoDB()
    ParseOBO(sys.argv[1], "%lineend%", MongoDBCollection)
    CloseMongoDB(MongoDBConnection)
